#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@author:  runyuanye
@contact: runyuanye@163.com
"""

import argparse
import os
import sys

sys.path.append('.')

"""
   对merge的人脸ID进行过滤,只保留每一个merge行的第一个faceid，其余的faceid都不要
"""


def read_features(file_list, feature_file, feat_dim):
    faceid_subfaceid_file_features = {}
    feat_file = open(feature_file, 'rb')
    file_count = 0
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read Feature Count: {}'.format(file_count))

            # if file_count == 10000:
            #     break

            # 20f10478-18a0-414c-a695-c43fb51077aa 11aa6f7c-b776-4e5b-8a16-fac29fbd5cb1 81c996ed4a8a743e046fa09025d69231.jpg 0
            lines = line.strip().split()

            faceid = lines[0]
            subfaceid = lines[1]
            # file_name = lines[2]
            # mode = int(lines[3])

            faceid_subfaceid = '{}_{}'.format(faceid, subfaceid)

            feature = feat_file.read(feat_dim * 4)

            file_features = faceid_subfaceid_file_features.get(faceid_subfaceid, None)
            if file_features is None:
                files = [line]
                features = [feature]
                faceid_subfaceid_file_features[faceid_subfaceid] = [files, features]
            else:
                file_features[0].append(line)
                file_features[1].append(feature)

        print('Read Feature Count: {}'.format(file_count))
    feat_file.close()
    return faceid_subfaceid_file_features


def read_merge_file(file_list):
    file_count = 0
    keep_faceid_subfaceids = set()
    skip_faceid_subfaceids = set()
    with open(file_list, 'r') as file:
        for line in file.readlines():
            file_count += 1
            if file_count % 10000 == 0:
                print('Read Feature Count: {}'.format(file_count))

            # if file_count == 10000:
            #     break

            # NONE:960362bc-3dec-4785-8c87-bd60052bb7f0_192a5bc2-ef6f-4ef5-b82f-9bc6cd0791eb,960362bc-3dec-4785-8c87-bd60052bb7f0_29685424-4d56-49fb-acb7-3416fba1bb15,960362bc-3dec-4785-8c87-bd60052bb7f0_dc9d79c5-cccf-4e21-8ece-e2a14f00325e
            lines = line.strip().split(':')
            faceids = lines[1].split(',')
            faceid = faceids[0]
            keep_faceid_subfaceids.add(faceid)
            skip_faceid_subfaceids.update(faceids[1:])
        print('Read Merge Count: {}'.format(file_count))
    return keep_faceid_subfaceids, skip_faceid_subfaceids


def main():
    parser = argparse.ArgumentParser(description="Replace Face ReID Image")
    parser.add_argument(
        "--file_list", default="/rootfs/media/kasim/Data/data/FaceImage/replace1/RepFaceRecSel.txt", help="path to image file list", type=str
    )
    parser.add_argument(
        "--feature_file", default="/rootfs/media/kasim/Data/data/FaceImage/replace1/FaceRecFeatSel.dat", help="path to feature file", type=str
    )
    parser.add_argument(
        "--merge_file", default="/rootfs/media/kasim/Data/data/FaceImage/replace1/MergeFace.txt", help="path to merge file", type=str
    )
    parser.add_argument(
        "--output_dir", default="/rootfs/media/kasim/Data/data/FaceImage/replace1", help="path to output", type=str
    )
    parser.add_argument(
        "--feature_dim", default=256, help="feature dim", type=int
    )
    parser.add_argument('--min_face_image_count', type=int, default=8, help='min face image count')
    args = parser.parse_args()

    output_dir = args.output_dir
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        os.system('chmod a+wr {}'.format(output_dir))

    faceid_subfaceid_file_features = read_features(args.file_list, args.feature_file, args.feature_dim)
    keep_faceid_subfaceids, skip_faceid_subfaceids = read_merge_file(args.merge_file)

    out_file_list_path = os.path.join(output_dir, 'Keep' + os.path.basename(args.file_list))
    out_feature_file_path = os.path.join(output_dir, 'Keep' + os.path.basename(args.feature_file))
    out_file_list = open(out_file_list_path, 'w')
    out_feature_file = open(out_feature_file_path, 'wb')

    face_count = 0
    face_image_count = 0
    for faceid_subfaceid, file_features in faceid_subfaceid_file_features.items():
        if (faceid_subfaceid in keep_faceid_subfaceids) or (faceid_subfaceid not in skip_faceid_subfaceids):
            files, features = file_features
            image_count = len(files)
            if image_count < args.min_face_image_count:
                continue
            face_count += 1
            face_image_count += image_count
            for file, feature in zip(files, features):
                out_file_list.write(file)
                out_feature_file.write(feature)
    out_feature_file.close()
    out_file_list.close()
    os.system('chmod a+wr {}'.format(out_file_list_path))
    os.system('chmod a+wr {}'.format(out_feature_file_path))
    print('Face Count: {}, Face Image Count: {}'.format(face_count, face_image_count))
    print('finish!')


if __name__ == '__main__':
    main()
